import tensorflow as tf
import os

os.environ["CUDA_VISIBLE_DEVICES"] = '-1'

"""https://arxiv.org/pdf/2005.09007.pdf """
"""U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection"""

class conv_block(tf.keras.layers.Layer):
    def __init__(self, out_ch=3, dirate=1, p=0.2, pool=True):
        super(conv_block, self).__init__()
        self.con_s1 = tf.keras.layers.Conv2D(out_ch,
                                             3,
                                             padding='same',
                                             dilation_rate=dirate)
        self.bn = tf.keras.layers.BatchNormalization(axis=3)
        self.relu = tf.keras.layers.ReLU(6.0)
        if pool:
            self.pool = tf.keras.layers.MaxPooling2D(pool_size=2, strides=2)
            self.dropout = tf.keras.layers.Dropout(p)
        self.pty = pool

    @tf.function
    def call(self, inputs, **kwargs):
        x = self.con_s1(inputs)
        x = self.bn(x)
        x = self.relu(x)
        if self.pty:
            p = self.pool(x)
            p = self.dropout(p)
            return x, p
        else:
            return x


class RSU7(tf.keras.layers.Layer):
    def __init__(self, out=64, M=32):
        super(RSU7, self).__init__()
        self.vis = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.ec1 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec2 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec3 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec4 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec5 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec6 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.ec7 = conv_block(out_ch=M, dirate=2, p=0.2, pool=False)

        # decode
        self.de1 = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.de2 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de3 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de4 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de5 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de6 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)

    @tf.function
    def call(self, inputs, **kwargs):
        inputs = self.vis(inputs)
        x1, p1 = self.ec1(inputs)
        x2, p2 = self.ec2(p1)
        x3, p3 = self.ec3(p2)
        x4, p4 = self.ec4(p3)
        x5, p5 = self.ec5(p4)
        x6 = self.ec6(p5)
        x7 = self.ec7(x6)

        # decode
        up6 = tf.concat([x6, x7], axis=-1)
        up6 = tf.keras.layers.UpSampling2D()(up6)
        d6 = self.de6(up6)

        up5 = tf.concat([x5, d6], axis=-1)
        up5 = tf.keras.layers.UpSampling2D()(up5)
        d5 = self.de5(up5)

        up4 = tf.concat([x4, d5], axis=-1)
        up4 = tf.keras.layers.UpSampling2D()(up4)
        d4 = self.de4(up4)

        up3 = tf.concat([x3, d4], axis=-1)
        up3 = tf.keras.layers.UpSampling2D()(up3)
        d3 = self.de3(up3)

        up2 = tf.concat([x2, d3], axis=-1)
        up2 = tf.keras.layers.UpSampling2D()(up2)
        d2 = self.de2(up2)

        up1 = tf.concat([x1, d2], axis=-1)
        d1 = self.de1(up1)

        return d1 + inputs


class RSU6(tf.keras.layers.Layer):
    def __init__(self, out=64, M=32):
        super(RSU6, self).__init__()
        self.vis = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.ec1 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec2 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec3 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec4 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec5 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.ec6 = conv_block(out_ch=M, dirate=2, p=0.2, pool=False)

        # decode
        self.de1 = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.de2 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de3 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de4 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de5 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)

    @tf.function
    def call(self, inputs, **kwargs):
        inputs = self.vis(inputs)
        x1, p1 = self.ec1(inputs)
        x2, p2 = self.ec2(p1)
        x3, p3 = self.ec3(p2)
        x4, p4 = self.ec4(p3)
        x5 = self.ec5(p4)
        x6 = self.ec6(x5)

        # decode
        up5 = tf.concat([x5, x6], axis=-1)
        up5 = tf.keras.layers.UpSampling2D()(up5)
        d5 = self.de5(up5)

        up4 = tf.concat([x4, d5], axis=-1)
        up4 = tf.keras.layers.UpSampling2D()(up4)
        d4 = self.de4(up4)

        up3 = tf.concat([x3, d4], axis=-1)
        up3 = tf.keras.layers.UpSampling2D()(up3)
        d3 = self.de3(up3)

        up2 = tf.concat([x2, d3], axis=-1)
        up2 = tf.keras.layers.UpSampling2D()(up2)
        d2 = self.de2(up2)

        up1 = tf.concat([x1, d2], axis=-1)
        d1 = self.de1(up1)

        return d1 + inputs


class RSU5(tf.keras.layers.Layer):
    def __init__(self, out=64, M=32):
        super(RSU5, self).__init__()
        self.vis = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.ec1 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec2 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec3 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec4 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.ec5 = conv_block(out_ch=M, dirate=2, p=0.2, pool=False)

        # decode
        self.de1 = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.de2 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de3 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de4 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)

    @tf.function
    def call(self, inputs, **kwargs):
        inputs = self.vis(inputs)
        x1, p1 = self.ec1(inputs)
        x2, p2 = self.ec2(p1)
        x3, p3 = self.ec3(p2)
        x4 = self.ec4(p3)
        x5 = self.ec5(x4)

        # decode
        up4 = tf.concat([x4, x5], axis=-1)
        up4 = tf.keras.layers.UpSampling2D()(up4)
        d4 = self.de4(up4)

        up3 = tf.concat([x3, d4], axis=-1)
        up3 = tf.keras.layers.UpSampling2D()(up3)
        d3 = self.de3(up3)

        up2 = tf.concat([x2, d3], axis=-1)
        up2 = tf.keras.layers.UpSampling2D()(up2)
        d2 = self.de2(up2)

        up1 = tf.concat([x1, d2], axis=-1)
        d1 = self.de1(up1)

        return d1 + inputs


class RSU4(tf.keras.layers.Layer):
    def __init__(self, out=64, M=32):
        super(RSU4, self).__init__()
        self.vis = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.ec1 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec2 = conv_block(out_ch=M, dirate=1, p=0.2)
        self.ec3 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.ec4 = conv_block(out_ch=M, dirate=2, p=0.2, pool=False)

        # decode
        self.de1 = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.de2 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)
        self.de3 = conv_block(out_ch=M, dirate=1, p=0.2, pool=False)

    @tf.function
    def call(self, inputs, **kwargs):
        inputs = self.vis(inputs)
        x1, p1 = self.ec1(inputs)
        x2, p2 = self.ec2(p1)
        x3 = self.ec3(p2)
        x4 = self.ec4(x3)

        # decode
        up3 = tf.concat([x3, x4], axis=-1)
        up3 = tf.keras.layers.UpSampling2D()(up3)
        d3 = self.de3(up3)

        up2 = tf.concat([x2, d3], axis=-1)
        up2 = tf.keras.layers.UpSampling2D()(up2)
        d2 = self.de2(up2)

        up1 = tf.concat([x1, d2], axis=-1)
        d1 = self.de1(up1)

        return d1 + inputs


class RSU4F(tf.keras.layers.Layer):
    def __init__(self, out=64, M=32):
        super(RSU4F, self).__init__()
        self.vis = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.ec1 = conv_block(out_ch=M, dirate=2, p=0.2, pool=False)
        self.ec2 = conv_block(out_ch=M, dirate=4, p=0.2, pool=False)
        self.ec3 = conv_block(out_ch=M, dirate=8, p=0.2, pool=False)

        # decode
        self.de1 = conv_block(out_ch=out, dirate=1, p=0.2, pool=False)
        self.de2 = conv_block(out_ch=M, dirate=2, p=0.2, pool=False)
        self.de3 = conv_block(out_ch=M, dirate=4, p=0.2, pool=False)

    @tf.function
    def call(self, inputs, **kwargs):
        inputs = self.vis(inputs)
        x1 = self.ec1(inputs)
        x2 = self.ec2(x1)
        x3 = self.ec3(x2)

        # decode
        up2 = tf.concat([x2, x3], axis=-1)
        d3 = self.de3(up2)
        up2 = tf.concat([x2, d3], axis=-1)
        d2 = self.de2(up2)
        up1 = tf.concat([x1, d2], axis=-1)
        d1 = self.de1(up1)

        return d1 + inputs


class U2net(tf.keras.Model):
    def __init__(self):
        super(U2net, self).__init__()
        self.EN_1 = RSU7(out=64, M=16)
        self.EN_2 = RSU6(out=64, M=16)
        self.EN_3 = RSU5(out=64, M=16)
        self.EN_4 = RSU4(out=64, M=16)
        self.EN_5 = RSU4F(out=64, M=16)
        self.EN_6 = RSU4F(out=64, M=16)

        self.DE_5 = RSU4F(out=64, M=16)
        self.DE_4 = RSU4(out=64, M=16)
        self.DE_3 = RSU5(out=64, M=16)
        self.DE_2 = RSU6(out=64, M=16)
        self.DE_1 = RSU7(out=64, M=16)

        self.pool = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))
        self.up = tf.keras.layers.UpSampling2D(size=(2, 2))

        self.up_block = []
        size = [1, 2, 4, 8, 16, 32]
        for s in size:
            SQ = tf.keras.Sequential()
            SQ.add(tf.keras.layers.Conv2D(1, 3, activation='sigmoid', padding='same'))
            SQ.add(tf.keras.layers.UpSampling2D(size=s))
            self.up_block.append(SQ)

        self.output_conv = tf.keras.layers.Conv2D(1, 1, activation='sigmoid', padding='same')

    @tf.function
    def call(self, inputs, training=None, mask=None):
        e1 = self.EN_1(inputs)
        p1 = self.pool(e1)
        e2 = self.EN_2(p1)
        p2 = self.pool(e2)
        e3 = self.EN_3(p2)
        p3 = self.pool(e3)
        e4 = self.EN_4(p3)
        p4 = self.pool(e4)
        e5 = self.EN_5(p4)
        p5 = self.pool(e5)
        e6 = self.EN_6(p5)

        s6 = self.up_block[5](e6)
        e6 = self.up(e6)
        d5 = self.DE_5(tf.concat([e5, e6], axis=-1))
        s5 = self.up_block[4](d5)

        d5up = self.up(d5)
        d4 = self.DE_4(tf.concat([e4, d5up], axis=-1))
        s4 = self.up_block[3](d4)

        d4up = self.up(d4)
        d3 = self.DE_3(tf.concat([e3, d4up], axis=-1))
        s3 = self.up_block[2](d3)

        d3up = self.up(d3)
        d2 = self.DE_2(tf.concat([e2, d3up], axis=-1))
        s2 = self.up_block[1](d2)

        d2up = self.up(d2)
        d1 = self.DE_1(tf.concat([e1, d2up], axis=-1))
        s1 = self.up_block[0](d1)

        s_fuse = tf.concat([s1, s2, s3, s4, s5, s6], axis=-1)
        s_fuse = self.output_conv(s_fuse)
        if training:
            return s1, s2, s3, s4, s5, s6, s_fuse
        else:
            return s_fuse


if __name__ == '__main__':
    model = U2net()
    raw = tf.io.read_file("../UNet/SKU130770.png")
    image = tf.image.decode_png(raw, channels=3)
    image = tf.image.resize(images=image, size=[224, 224])
    image = image / 127.5 - 1
    image = tf.expand_dims(image, 0)

    res = model(image)
    print(res.shape)
